
import numpy as np
import open3d as o3d

def visualize_point_cloud_with_labels(original_path, label_path):
    # 1. 加载数据
    original_data = np.load(original_path)  # 原始点云 [x,y,z,r,g,b,...]
    pred_labels = np.load(label_path)       # 预测标签 [N,]

    # 2. 检查数据格式和一致性
    assert len(original_data) == len(pred_labels), "点云和标签数量不匹配!"
    print("原始点云形状:", original_data.shape)
    print("标签形状:", pred_labels.shape)

    # 3. 提取坐标和原始颜色（假设原始点云有RGB）
    points = original_data[:, :3]
    original_colors = original_data[:, 3:6] / 255.0 if original_data.shape[1] >= 6 else None

    # 4. 为预测标签生成伪彩色
    label_colors = generate_label_colors(pred_labels)

    # 5. 创建Open3D点云对象（原始点云和带标签点云）
    pcd_original = o3d.geometry.PointCloud()
    pcd_original.points = o3d.utility.Vector3dVector(points)
    if original_colors is not None:
        pcd_original.colors = o3d.utility.Vector3dVector(original_colors)

    pcd_label = o3d.geometry.PointCloud()
    pcd_label.points = o3d.utility.Vector3dVector(points)
    pcd_label.colors = o3d.utility.Vector3dVector(label_colors)

    # 6. 可选：下采样（如果点云过大）
    if len(points) > 100000:
        voxel_size = 0.03
        pcd_original = pcd_original.voxel_down_sample(voxel_size)
        pcd_label = pcd_label.voxel_down_sample(voxel_size)

    # 7. 可视化（分窗口显示原始点云和标签点云）
    visualize_side_by_side(pcd_original, pcd_label, "原始点云 (RGB)", "预测标签 (伪彩色)")

def generate_label_colors(labels):
    """为标签生成鲜艳的伪彩色（固定颜色或随机颜色）"""
    unique_labels = np.unique(labels)
    # 固定颜色映射（例如：0=红色，1=绿色，2=蓝色...）
    color_map = {
        0: [1, 0, 0],   # 红
        1: [0, 1, 0],   # 绿
        2: [0, 0, 1],   # 蓝
        3: [1, 1, 0],   # 黄
        4: [1, 0, 1],   # 紫
        5: [0, 1, 1],   # 青
    }
    # 对于未定义的标签，随机生成颜色
    for label in unique_labels:
        if label not in color_map:
            color_map[label] = np.random.rand(3)
    return np.array([color_map[label] for label in labels])

def visualize_side_by_side(pcd1, pcd2, title1, title2):
    """左右分窗口显示两个点云"""
    # 窗口1：原始点云
    vis1 = o3d.visualization.Visualizer()
    vis1.create_window(window_name=title1, width=800, height=600, left=0)
    vis1.add_geometry(pcd1)
    opt1 = vis1.get_render_option()
    opt1.point_size = 2.0
    opt1.background_color = np.array([0.1, 0.1, 0.1])

    # 窗口2：标签点云
    vis2 = o3d.visualization.Visualizer()
    vis2.create_window(window_name=title2, width=800, height=600, left=800)
    vis2.add_geometry(pcd2)
    opt2 = vis2.get_render_option()
    opt2.point_size = 2.0
    opt2.background_color = np.array([0.1, 0.1, 0.1])

    # 同时运行两个窗口
    while True:
        vis1.update_geometry(pcd1)
        vis2.update_geometry(pcd2)
        if not vis1.poll_events() or not vis2.poll_events():
            break
        vis1.update_renderer()
        vis2.update_renderer()

    vis1.destroy_window()
    vis2.destroy_window()


# 调用函数
# visualize_point_cloud_with_labels(
#     "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_office_26.npy",
#     "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-office_26_pred.npy"
# )

# visualize_point_cloud_with_labels(
#     "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_hallway_8.npy",
#     "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-hallway_8_pred.npy"
# )

# visualize_point_cloud_with_labels(
#     "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_hallway_10.npy",
#     "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-hallway_10_pred.npy"
# )


# visualize_point_cloud_with_labels(
#     "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_hallway_12.npy",
#     "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-hallway_12_pred.npy"
# )

# visualize_point_cloud_with_labels(
#     "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_office_11.npy",
#     "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-office_11_pred.npy"
# )

visualize_point_cloud_with_labels(
    "E:\\data\\3DpointCloud\\s3dis\\trainval_fullarea\\Area_5_storage_4.npy",
    "D:\\resource\code\\202504\\3dponits\\pred\\Area_5-storage_4_pred.npy"
)
